import numpy as np
import torch
import torch.backends.cudnn as cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import parser
from datalist import VideoDataset2
from model import C3D

best_acc = 0


class train(object):
    def __init__(self):
        self.args = parser.parse_args()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)  # 为当前GPU设置随机种子
        else:
            torch.manual_seed(self.args.seed)  # 为CPU设置种子用于生成随机数，以使得结果是确定的
        self.device = torch.device("cuda" if use_cuda else "cpu")
        val_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}

        '''
        构造DataLoader
        '''

        self.test_dataloader = DataLoader(VideoDataset2(self.args.train_dir), batch_size=self.args.test_batch_size,
                                          shuffle=True, **val_kwargs)

        self.model = C3D(2).to(self.device)

        if self.args.resume:
            print("load the weight from pretrained-weight file")
            model_dict = self.model.state_dict()
            pretrained_dict = torch.load(self.args.pretrained_weight, map_location=self.device)["model_state_dict"]
            pretrained_dict = {k[7:]: v for k, v in pretrained_dict.items() if
                               np.shape(model_dict[k[7:]]) == np.shape(v)}
            model_dict.update(pretrained_dict)
            self.model.load_state_dict(model_dict)
            print("Finished to load the weight")

        if use_cuda:
            self.model = torch.nn.DataParallel(self.model, device_ids=range(torch.cuda.device_count()))
            cudnn.benchmark = True

        '''
        构造loss目标函数
        选择优化器
        学习率变化选择
        '''

        # self.criterion = nn.CrossEntropyLoss()
        # self.optimizer = optim.SGD(self.model.parameters(), lr=self.args.lr, momentum=self.args.momentum)
        # self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=5, eta_min=1e-5)

        for epoch in range(1, 2):
            self.test(epoch)

        torch.cuda.empty_cache()
        print("finish model training")

    def test(self, epoch):

        global best_acc

        self.model.eval()
        test_loss = 0
        correct = torch.zeros(1).squeeze().cuda()
        total = torch.zeros(1).squeeze().cuda()

        average_loss = []

        pbar = tqdm(self.test_dataloader,
                    desc=f'Test Epoch {epoch}/{self.args.epoches}',
                    mininterval=0.3)
        for data, target in pbar:
            data, target = data.to(self.device), target.to(self.device)
            with torch.no_grad():
                output = self.model(data)
            # average_loss.append(self.criterion(output, target).item())
            # test_loss += self.criterion(output, target).item()  # sum up batch loss
            pred = torch.argmax(output, 1)
            target = torch.argmax(target, 1)

            correct += (pred == target).sum().float()
            total += len(target)
            predict_acc = correct / total
            pbar.set_description(
                f'Test Epoch:{epoch}/{self.args.epoches} acc:{predict_acc}')


if __name__ == "__main__":
    train = train()
